""" 
假如60日均线小于120日均线时，并且当5日均线与30均线金叉-买总资金的10%。
假如60日均线大于120日均线, 并且当5日均线与30日均线死叉卖10%
"""
from datetime import datetime
import backtrader as bt
import matplotlib.pyplot as plt
import akshare as ak
import pandas as pd

plt.rcParams["font.sans-serif"] = ["SimHei"]
plt.rcParams["axes.unicode_minus"] = False

# 获取股票后复权数据
stock_hfq_df = ak.stock_zh_a_hist(symbol="002508", adjust="qfq").iloc[:, :6]
# 处理字段命名
stock_hfq_df.columns = ['date', 'open', 'close', 'high', 'low', 'volume']
# 将date设为索引
stock_hfq_df.index = pd.to_datetime(stock_hfq_df['date'])


class MyStrategy(bt.Strategy):
    """
    主策略程序
    """
    params = (
        ('fast_ma', 5),  # 快速均线周期
        ('slow_ma', 20),  # 慢速均线周期
        ('long_ma1', 60),  # 长期均线1周期
        ('long_ma2', 120),  # 长期均线2周期
        ('trade_percentage', 0.1),  # 交易资金比例
    )
    
    def __init__(self):
        self.data_close = self.datas[0].close
        self.fast_sma = bt.indicators.SimpleMovingAverage(self.data.close, period=self.params.fast_ma)
        self.slow_sma = bt.indicators.SimpleMovingAverage(self.data.close, period=self.params.slow_ma)
        self.long_sma1 = bt.indicators.SimpleMovingAverage(self.data.close, period=self.params.long_ma1)
        self.long_sma2 = bt.indicators.SimpleMovingAverage(self.data.close, period=self.params.long_ma2)

    def next(self):
        current_date = self.data.datetime.date(0)
        size = 100
        if self.long_sma1[0] > self.long_sma2[0]:
            #if not self.position:  # 没有仓位时检查买入条件
            if self.fast_sma[0] > self.slow_sma[0] and self.fast_sma[-1] <= self.slow_sma[-1]:
                
                cash_available = self.broker.getcash()  # 获取可用资金                
                price = self.data.close[0]  # 当前价格
                amount_to_invest = self.broker.cash * self.params.trade_percentage
                #self.buy(size=amount_to_invest / self.data.close[0])
                self.buy(size=size)  # 执行买入100股的操作
                print(f"{current_date}: 当前价格{price}, 买入{amount_to_invest}")
            else:
                print(f"{current_date}: 余额不足，无法购买。")
        elif self.long_sma1[0] <= self.long_sma2[0]:  # 有仓位时检查卖出条件
            if self.fast_sma[0] < self.slow_sma[0]and self.fast_sma[-1] >= self.slow_sma[-1]:
                price = self.data.close[0]  # 当前价格
                self.sell(size=size)
                print(f"{current_date}: 当前价格{price}, 卖出{self.position.size}")
            else:
                if self.position.size <= 0:
                    print(f"{current_date}: 余额不足，无法卖出。")                

# 初始化回测系统
cerebro = bt.Cerebro()

# 设置回测时间范围
start_date = datetime(2000, 4, 3)
end_date = datetime(2024, 5, 22)

# 加载数据到Cerebro
data = bt.feeds.PandasData(dataname=stock_hfq_df, fromdate=start_date, todate=end_date)
cerebro.adddata(data)

# 加载策略到Cerebro
cerebro.addstrategy(MyStrategy)

# 设置初始资金和手续费
start_cash = 1000000
cerebro.broker.setcash(start_cash)
cerebro.broker.setcommission(commission=0.002)

# 运行回测
cerebro.run()

# 获取回测结果并打印
port_value = cerebro.broker.getvalue()
pnl = port_value - start_cash
print(f"初始资金: {start_cash}\n回测期间：{start_date.strftime('%Y%m%d')}:{end_date.strftime('%Y%m%d')}")
print(f"总资金: {round(port_value, 2)}")
print(f"净收益: {round(pnl, 2)}")

# 绘制图表
cerebro.plot(style='candlestick')